package redisx

import (
	"context"
	"github.com/redis/go-redis/v9"
	"go.opentelemetry.io/otel"
	"go.opentelemetry.io/otel/attribute"
	"go.opentelemetry.io/otel/codes"
	"go.opentelemetry.io/otel/trace"
	"net"
)

type OpenTelemetryHook struct {
	tracer trace.Tracer
}

func NewOpenTelemetryHook() *OpenTelemetryHook {
	return &OpenTelemetryHook{
		tracer: otel.GetTracerProvider().Tracer("Redis Tracing"),
	}
}

func (o *OpenTelemetryHook) DialHook(next redis.DialHook) redis.DialHook {
	return func(ctx context.Context, network, addr string) (net.Conn, error) {
		return next(ctx, network, addr)
	}
}

func (o *OpenTelemetryHook) ProcessHook(next redis.ProcessHook) redis.ProcessHook {
	return func(ctx context.Context, cmd redis.Cmder) error {
		ctx, span := o.tracer.Start(ctx, "Redis:"+cmd.Name(), trace.WithSpanKind(trace.SpanKindClient))
		defer span.End()
		span.SetAttributes(attribute.String("name", cmd.Name()))
		span.SetAttributes(attribute.String("string", cmd.String()))
		err := next(ctx, cmd)
		if err != nil {
			span.RecordError(err)
			span.SetStatus(codes.Error, err.Error())
		}
		return err
	}
}

func (o *OpenTelemetryHook) ProcessPipelineHook(next redis.ProcessPipelineHook) redis.ProcessPipelineHook {
	return func(ctx context.Context, cmds []redis.Cmder) error {
		return next(ctx, cmds)
	}
}
